/*
 * Copyright (c) 2016-2021, Hangzhou Qulian Technology Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "keccak256.h"

#define BLOCK_SIZE     ((1600 - 256 * 2) / 8)
#define ROTL64(qword, n) ((qword) << (n) ^ ((qword) >> (64 - (n))))
#define le2me_64(x) (x)
#define IS_ALIGNED_64(p) (0 == (7 & ((const OhByteType *)(p) - (const OhByteType *)0)))
#define me64_to_le_str(to, from, length) memcpy((to), (from), (length))

// -----------------------
// sub-operation process
// THETA RHO PI CHI ROUND
// -----------------------
#define TYPE_ROUND_INFO      0
#define TYPE_PI_TRANSFORM   24
#define TYPE_RHO_TRANSFORM  48
const OhUint8Type constants[] = {
        1,  26,  94, 112,  31,  33, 121,  85,  14,  12,  53,  38,  63,  79,  93,  83,  82,  72,  22, 102, 121,  88,  33, 116,
        1,   6,   9,  22,  14,  20,   2,  12,  13,  19,  23,  15,   4,  24,  21,   8,  16,   5,   3,  18,  17,  11,   7,  10,
        1,  62,  28,  27,  36,  44,   6,  55,  20,   3,  10,  43,  25,  39,  41,  45,  15,  21,   8,  18,   2,  61,  56,  14,
};
OhUint8Type getConstant(OhUint8Type type, OhUint8Type index) {
    return constants[type + index];
}

// Keccak theta() transformation
static OhVoidType keccak_theta(OhUint64Type *A) {

    OhUint64Type C[5], D[5];
    OhUint8Type i, j;
    for (i = 0; i < 5; i++) {
        C[i] = A[i];
        for (j = 5; j < 25; j += 5) { C[i] ^= A[i + j]; }
    }
    for (i = 0; i < 5; i++) {
        D[i] = ROTL64(C[(i + 1) % 5], 1) ^ C[(i + 4) % 5];
    }
    for (i = 0; i < 5; i++) {
        for (j = 0; j < 25; j += 5) { A[i + j] ^= D[i]; }
    }
}

// Keccak pi() transformation
static OhVoidType keccak_pi(OhUint64Type *A) {

    OhUint64Type A1 = A[1];
    OhUint8Type i;
    for (i = 1; i < 24; i++) {
        A[getConstant(TYPE_PI_TRANSFORM, i - 1)] = A[getConstant(TYPE_PI_TRANSFORM, i)];
    }
    A[10] = A1;
}

// Keccak chi() transformation
static OhVoidType keccak_chi(OhUint64Type *A) {
    
    OhUint8Type i;
    for (i = 0; i < 25; i += 5) {
        OhUint64Type A0 = A[0 + i], A1 = A[1 + i];
        A[0 + i] ^= ~A1 & A[2 + i];
        A[1 + i] ^= ~A[2 + i] & A[3 + i];
        A[2 + i] ^= ~A[3 + i] & A[4 + i];
        A[3 + i] ^= ~A[4 + i] & A0;
        A[4 + i] ^= ~A0 & A1;
    }
}

// Keccak l() transformation
static OhUint64Type get_round_constant(OhUint8Type round) {

    OhUint64Type result = 0;
    OhUint8Type roundInfo = getConstant(TYPE_ROUND_INFO, round);
    if (roundInfo & (1 << 6)) { result |= ((OhUint64Type) 1 << 63); }
    if (roundInfo & (1 << 5)) { result |= ((OhUint64Type) 1 << 31); }
    if (roundInfo & (1 << 4)) { result |= ((OhUint64Type) 1 << 15); }
    if (roundInfo & (1 << 3)) { result |= ((OhUint64Type) 1 <<  7); }
    if (roundInfo & (1 << 2)) { result |= ((OhUint64Type) 1 <<  3); }
    if (roundInfo & (1 << 1)) { result |= ((OhUint64Type) 1 <<  1); }
    if (roundInfo & (1 << 0)) { result |= ((OhUint64Type) 1 <<  0); }
    return result;
}

// make a permutation of the hash
static OhVoidType sha3_permutation(OhUint64Type *state) {
    OhUint8Type round, i;
    for (round = 0; round < 24; round++) {
        keccak_theta(state);                  // 1
        for (i = 1; i < 25; i++) { // 2 apply Keccak rho() transformation
            state[i] = ROTL64(state[i], getConstant(TYPE_RHO_TRANSFORM, i - 1));
        }
        keccak_pi(state);                     // 3
        keccak_chi(state);                    // 4
        *state ^= get_round_constant(round);  // 5 apply iota(state, round)
    }
}

// The core transformation. Process the specified block of data
static OhVoidType sha3_process_block(OhUint64Type hash[25], const OhUint64Type *block) {
    OhUint8Type i;
    for (i = 0; i < 17; i++) {
        hash[i] ^= le2me_64(block[i]);
    }
    sha3_permutation(hash);
}

// -----------------------
// Calculate message hash
// -----------------------
// define sha3 content
typedef struct sha3_ctx_t {
    OhUint64Type hash[SHA3_MAX_PERMUTATION_SIZE];    // 1600 bits algorithm hashing state
    OhUint64Type message[SHA3_MAX_RATE_IN_QWORDS];   // 1536-bit buffer for leftovers
    OhUint16Type rest;                               // count of bytes in the message[] buffer
} sha3_ctx_t;

// sha3 buffer init
// ctx   : the algorithm context containing current hashing state
OhVoidType keccak_init(sha3_ctx_t *ctx) {
    memset(ctx, 0, sizeof(sha3_ctx_t));
}

// Calculate message hash. Can be called repeatedly with chunks of the message to be hashed.
// ctx : the algorithm context containing current hashing state
// msg : message chunk
OhVoidType keccak_update(sha3_ctx_t *ctx, const OhByteType *msg, OhUint16Type size) {

    OhUint16Type idx = ctx->rest;
    ctx->rest = (ctx->rest + size) % BLOCK_SIZE;
    /* fill partial block */
    if (idx) {
        OhUint16Type left = BLOCK_SIZE - idx;
        memcpy((OhByteType *) ctx->message + idx, msg, (size < left ? size : left));
        if (size < left) return;
        /* process partial block */
        sha3_process_block(ctx->hash, ctx->message);
        msg += left;
        size -= left;
    }
    while (size >= BLOCK_SIZE) {
        OhUint64Type *aligned_message_block;
        if (IS_ALIGNED_64(msg)) {
            // the most common case is processing of an already aligned message without copying it
            aligned_message_block = (OhUint64Type *) (void *) msg;
        } else {
            memcpy(ctx->message, msg, BLOCK_SIZE);
            aligned_message_block = ctx->message;
        }
        sha3_process_block(ctx->hash, aligned_message_block);
        msg  += BLOCK_SIZE;
        size -= BLOCK_SIZE;
    }
    if (size) {
        memcpy(ctx->message, msg, size); /* save leftovers */
    }
}

// Store calculated hash into the given array.
// ctx   : the algorithm context containing current hashing state
// param : result calculated hash in binary form
OhVoidType keccak_final(sha3_ctx_t *ctx, OhByteType *result) {

    OhUint16Type digest_length = 100 - BLOCK_SIZE / 2;
    /* clear the rest of the data queue */
    memset((OhByteType *) ctx->message + ctx->rest, 0, BLOCK_SIZE - ctx->rest);
    ((OhByteType *) ctx->message)[ctx->rest     ] |= 0x01;
    ((OhByteType *) ctx->message)[BLOCK_SIZE - 1] |= 0x80;
    /* process final block */
    sha3_process_block(ctx->hash, ctx->message);
    if (result) {
        me64_to_le_str(result, ctx->hash, digest_length);
    }
}

// keccak256 is the package of keccak_init keccak_update keccak_final.
// msg     : the raw data need to hash
// size    : data length
// hashbuf : store result after keccak256
OhVoidType keccak256(const OhByteType *msg, OhUint32Type size, OhByteType **hashbuf) {
    sha3_ctx_t ctx;
    keccak_init(&ctx);
    keccak_update(&ctx, msg, size);
    keccak_final(&ctx, *hashbuf);
}

// -----------------------
// Test function
// -----------------------
// print result (hex data = 32 byte)
OhVoidType print_keccak_data(OhByteType *out){
    OhByteType buf[65];
    OhUint8Type i;
    for(i = 0; i < 32; i ++){
        buf[i*2]   = ((out[i] & 0xf0) >> 4) <=9 ?  ((out[i] & 0xf0) >> 4) + '0' : ((out[i] & 0xf0) >> 4) + 'a' - 10;
        buf[i*2+1] = ((out[i] & 0x0f)     ) <=9 ?  ((out[i] & 0x0f)     ) + '0' : ((out[i] & 0x0f)     ) + 'a' - 10;
    }
    buf[64] = '\0';
    log_info("[keccak data] %s",buf);
}
